#include "benchmark.hpp"
#include "probe.hpp"
#include "settingsFunctions.hpp"
#include "bc_base_all.hpp"
#include "bc_swe_all.hpp"
#include "timing.hpp"
#include "blockwiseIntegrated.hpp"
#include "loggerFunctions.hpp"
#include "infiltration.hpp"

#include "vtkImport.hpp"
#include "vtkExport.hpp"

#include "timeMarchingLoop.hpp"
#include "exporter.hpp"

namespace hms
{

template<typename MeshType>
void setBenchCase_lateralFlow(
	SWESolver<MeshType>& solver,
	const MeshGenSettings& meshGenSettings
){
	Field<MeshType>
		& d { solver.d() },
		& v { solver.v() },
		& z { solver.z() };

	d.values() = 1;
	v.values().colwise() = Array2s{1, 0};

	if ( !meshGenSettings.addBoundaryPatches ) return;

	if ( !meshGenSettings.useStandardBoundaryPatches )
		throw MeshError(
			"Benchmark cases are only set up for standard boundary patches!"
		);
	
	const MeshType& mesh { solver.mesh() };

	const BoundaryPatch
		& bottom { mesh.boundaryPatch("bottom") },
		& top    { mesh.boundaryPatch("top") },
		& right  { mesh.boundaryPatch("right") },
		& left   { mesh.boundaryPatch("left") };

	scalar vGhost {1.0};
	[[maybe_unused]] scalar dGhost {1.0};

	d.setBoundaryCondition( OrthogonalFlowVelocity_waterdepth{d, left, v, vGhost} );
	v.setBoundaryCondition( OrthogonalFlowVelocity_velocity  {v, left,    vGhost} );

	d.setBoundaryCondition( FreeOutflow_waterdepth{d, right, false} );
	v.setBoundaryCondition( FreeOutflow_velocity  {v, right} );

	d.setBoundaryCondition( ZeroGradient {d, top} );
	v.setBoundaryCondition( Slip_velocity{v, top} );

	d.setBoundaryCondition( ZeroGradient {d, bottom} );
	v.setBoundaryCondition( Slip_velocity{v, bottom} );

	z.setBoundaryCondition( StaticValue{z, left  , false} );
	z.setBoundaryCondition( StaticValue{z, right , false} );
	z.setBoundaryCondition( StaticValue{z, bottom, false} );
	z.setBoundaryCondition( StaticValue{z, top   , false} );
}

template<typename MeshType>
void setBenchCase_diagonalFlow(
	SWESolver<MeshType>& solver,
	const MeshGenSettings& meshGenSettings
){
	Field<MeshType>
		& d { solver.d() },
		& v { solver.v() },
		& z { solver.z() };

	d.values() = 1;
	v.values().colwise() = Array2s{1, 1};

	if ( !meshGenSettings.addBoundaryPatches ) return;
	
	const MeshType& mesh { solver.mesh() };

	const BoundaryPatch
		& bottom { mesh.boundaryPatch("bottom") },
		& top    { mesh.boundaryPatch("top") },
		& right  { mesh.boundaryPatch("right") },
		& left   { mesh.boundaryPatch("left") };

	scalar vGhost {1.0};
	[[maybe_unused]] scalar dGhost {1.0};

	d.setBoundaryCondition( OrthogonalFlowVelocity_waterdepth{d, left, v, vGhost} );
	v.setBoundaryCondition( OrthogonalFlowVelocity_velocity  {v, left,    vGhost} );

	d.setBoundaryCondition( OrthogonalFlowVelocity_waterdepth{d, bottom, v, vGhost} );
	v.setBoundaryCondition( OrthogonalFlowVelocity_velocity  {v, bottom,    vGhost} );

	d.setBoundaryCondition( FreeOutflow_waterdepth{d, right, false} );
	v.setBoundaryCondition( FreeOutflow_velocity  {v, right} );

	d.setBoundaryCondition( FreeOutflow_waterdepth{d, top, false} );
	v.setBoundaryCondition( FreeOutflow_velocity  {v, top} );

	z.setBoundaryCondition( StaticValue{z, left  , false} );
	z.setBoundaryCondition( StaticValue{z, right , false} );
	z.setBoundaryCondition( StaticValue{z, bottom, false} );
	z.setBoundaryCondition( StaticValue{z, top   , false} );
}

template<typename MeshType>
void setBenchCase_dambreak(
	SWESolver<MeshType>& solver,
	const BenchmarkSettings& benchSettings,
	const MeshGenSettings& meshGenSettings
){
	using C = typename BenchmarkSettings::BenchCase;
	scalar dHigh {4}, dLow {
		benchSettings.benchCase == C::dambreak_wet ? 1.0 : 0.0
	};

	Field<MeshType>
		& d { solver.d() },
		& v { solver.v() },
		& z { solver.z() };

	v.values().colwise() = Array2s{0, 0};
	
	const MeshType& mesh { solver.mesh() };

	scalar middleX {
		meshGenSettings.originX + 
		0.5 * meshGenSettings.nCellsX * meshGenSettings.dX
	};

	mesh.forEachCell( [&](const auto& cell){
		d.values()( 0, cell.index() ) = 
			cell.centroid().x() < middleX ? dHigh : dLow;
	});

	if ( !meshGenSettings.addBoundaryPatches ) return;

	const BoundaryPatch
		& bottom { mesh.boundaryPatch("bottom") },
		& top    { mesh.boundaryPatch("top") },
		& right  { mesh.boundaryPatch("right") },
		& left   { mesh.boundaryPatch("left") };

	d.setBoundaryCondition( FreeOutflow_waterdepth{d, left, false} );
	v.setBoundaryCondition( FreeOutflow_velocity  {v, left} );

	d.setBoundaryCondition( FreeOutflow_waterdepth{d, right, false} );
	v.setBoundaryCondition( FreeOutflow_velocity  {v, right} );

	d.setBoundaryCondition( ZeroGradient {d, top} );
	v.setBoundaryCondition( Slip_velocity{v, top} );

	d.setBoundaryCondition( ZeroGradient {d, bottom} );
	v.setBoundaryCondition( Slip_velocity{v, bottom} );

	z.setBoundaryCondition( StaticValue{z, left  , false} );
	z.setBoundaryCondition( StaticValue{z, right , false} );
	z.setBoundaryCondition( StaticValue{z, bottom, false} );
	z.setBoundaryCondition( StaticValue{z, top   , false} );
}

template<typename MeshType>
void setBenchCase(
	SWESolver<MeshType>& solver,
	const BenchmarkSettings& benchSettings,
	const MeshGenSettings& meshGenSettings
){
	if ( !meshGenSettings.useStandardBoundaryPatches )
		throw MeshError(
			"Benchmark cases are only set up for standard boundary patches!"
		);

	using C = typename BenchmarkSettings::BenchCase;
	switch(benchSettings.benchCase){
		case C::lateralFlow:
			setBenchCase_lateralFlow(solver, meshGenSettings);
			return;
		case C::diagonalFlow:
			setBenchCase_diagonalFlow(solver, meshGenSettings);
			return;
		case C::dambreak_dry: [[fallthrough]];
		case C::dambreak_wet:
			setBenchCase_dambreak(solver, benchSettings, meshGenSettings);
			return;
		default:
			return;
	}
}

void printPrepTimings( const PrepTimings& timings ){
	printTimings<PreprocessingStep>( timings,
		"Total preprocessing time",
		{{
			"mesh creation",
			"solver creation",
			"solver initialisation",
			"setting fields",
			"output"
		}}
	);
}

void runBenchmark( Settings& settings ){
	#define BENCHMARK_TIME_PREPROCESSING_STEPS 
	#ifdef BENCHMARK_TIME_PREPROCESSING_STEPS
		using PS = PreprocessingStep;
		Timer prepTimer;
		PrepTimings prepTimings { PrepTimings::Zero() };
		#define TIME_PREP_STEP(STEP) \
			setTiming( prepTimings, PS::STEP, prepTimer)
		#define PRINT_PREP_TIMINGS \
			printPrepTimings( prepTimings )
	#else
		#define TIME_PREP_STEP(STEP)
		#define PRINT_PREP_TIMINGS
	#endif

	MeshVariant meshVar { generateMesh( settings.meshGen() ) };
	setBlocksize( meshVar, settings.swe() );
	TIME_PREP_STEP(meshCreation);

	std::visit( [&](const auto& mesh){
		log::logger& logger { mesh.logger() };

		Time time { std::move( settings.time() ) };

		SWESolver solver {mesh, time};

		solver.settings() = std::move( settings.swe() );
		setBlocksizes(solver);
		setInfiltrationCallable(solver);
		TIME_PREP_STEP(solverCreation);
	
		setBenchCase( solver, settings.benchmark(), settings.meshGen() );
		TIME_PREP_STEP(benchSettings);

		log::printSection(logger, "Case setup");

		logger.info(
			"Running benchmark case \"{}\", using\n{}"
			"{} x {} cells with size {}m x {}m",
			magic_enum::enum_name(settings.benchmark().benchCase),
			indent(),
			settings.meshGen().nCellsX,
			settings.meshGen().nCellsY,
			settings.meshGen().dX,
			settings.meshGen().dY
		);
		printTimeSettings( time, logger );

		mesh.printInfo();
		// TIME_PREP_STEP(output);

		printSolverSettings(solver);
		// TIME_PREP_STEP(output);

		printFieldProperties(solver);
		TIME_PREP_STEP(output);

		/* make sure everything is in order */
		solver.initialise();
		TIME_PREP_STEP(solverInit);
		PRINT_PREP_TIMINGS;

		if (settings.dryRun)
			return;

		tinyxml2::XMLDocument bindingDoc;

		const bool& exportResults { settings.benchmark().exportBenchmarkResults };

		if (exportResults){
			vtk::createBindingDoc( &bindingDoc );
		}
	
		log::printSection(logger, "Running benchmark");

		timeMarchingLoop(solver, time, Exporter{
			exportResults, settings.caseDir(), settings.caseName(),
			solver, bindingDoc, settings.bindingFilePath(),
			std::move( settings.probes() )
		} );

		log::printSection(logger, "Benchmark finished");
		
	}, meshVar );
}
	
} // namespace hms
